Summary of Changes

This PR adds support for FAIRChem batched inference via the InferenceBatcher from fairchem.core.calculate. This enables efficient GPU batching when running the same calculation (e.g., relaxations, static calculations) on many structures concurrently.

Key additions:

1. New map_partition_fairchembatch function in src/quacc/wflow_tools/job_patterns.py:

   • A @job that processes a single partition (flat list) of atoms using FAIRChem batched inference
   • Submits calculations concurrently via threads while Ray Serve batches GPU requests

2. New map_partitioned_lists_fairchembatch function in src/quacc/wflow_tools/job_patterns.py:

   • Dispatches partitioned data (list of lists) to map_partition_fairchembatch jobs
   • Mirrors the existing map_partitioned_lists / map_partition pattern

3. New get_inference_batcher function in src/quacc/recipes/mlp/_base.py:

   • Creates and caches InferenceBatcher instances keyed by model configuration
   • Includes shutdown_inference_batchers() for cleanup

4. Updated pick_calculator in src/quacc/recipes/mlp/_base.py:

   • Now accepts predict_unit kwarg to use batched inference
   • Bypasses caching when predict_unit is provided for thread-safety during concurrent execution

Example usage:

from quacc.recipes.mlp.core import relax_job
from quacc.wflow_tools.job_patterns import (
    map_partitioned_lists_fairchembatch,
    partition,
    unpartition,
)

# A very large array of possible atoms objects or structures to apply a job to
my_atoms_list = ...

# Partition atoms into batches for distributed execution
partitioned_atoms = partition(my_atoms_list, num_partitions = 10)

# Run batched FAIRChem relaxations on each partition
# Each partition runs on a separate worker with GPU batching
results_partitioned = map_partitioned_lists_fairchembatch(
        relax_job,
        atoms_list=partitioned_atoms,
        fairchem_model="uma-s-1",
        task_name="omat",
)

# Recombine results into a single list
all_results = unpartition(results_partitioned)

Tests:

• Comprehensive test suite in test_fairchem_batch.py]
• Includes end-to-end test with 10 copper structures
• Includes performance test verifying batched inference speedup vs sequential

Requirements

• My PR is focused on a single feature addition or bugfix.
• My PR has relevant, comprehensive unit tests.
• My PR is on a custom branch (i.e. is not named main).